package pers.qianyu.month_202103.date_20210317;

import java.util.HashMap;
import java.util.Map;

/**
 * 208. 实现 Trie (前缀树)
 * https://leetcode-cn.com/problems/implement-trie-prefix-tree/
 * - 字典树的另一种实现方式
 *
 * @author mizzle rain
 * @date 2021-03-17 10:40
 */
public class Trie {
    private final TrieNode root;

    /**
     * Initialize your data structure here.
     */
    public Trie() {
        root = new TrieNode();
    }

    /**
     * Inserts a word into the trie.
     */
    public void insert(String word) {
        int n = word.length();
        char[] arr = word.toCharArray();
        TrieNode p = root;
        for (int i = 0; i < n; i++) {
            if (!p.containsKey(arr[i])) {
                p.put(arr[i], new TrieNode());
            }
            p = p.get(arr[i]);
        }
        p.setEnd(true);
    }

    /**
     * Returns if the word is in the trie.
     */
    public boolean search(String word) {
        TrieNode node = searchPrefix(word);
        return node != null && node.isEnd();
    }

    /**
     * Returns if there is any word in the trie that starts with the given prefix.
     */
    public boolean startsWith(String prefix) {
        TrieNode node = searchPrefix(prefix);
        return node != null;
    }

    private TrieNode searchPrefix(String word) {
        TrieNode p = root;
        int n = word.length();
        char[] arr = word.toCharArray();
        for (int i = 0; i < n; i++) {
            if (!p.containsKey(arr[i])) {
                return null;
            }
            p = p.get(arr[i]);
        }
        return p;
    }

    static class TrieNode {
        private boolean isEnd;
        private final Map<Character, TrieNode> links;

        public TrieNode() {
            links = new HashMap<>();
        }

        public boolean containsKey(Character ch) {
            return links.containsKey(ch);
        }

        public void put(Character ch, TrieNode node) {
            links.put(ch, node);
        }

        public TrieNode get(Character ch) {
            return links.get(ch);
        }

        public boolean isEnd() {
            return isEnd;
        }

        public void setEnd(boolean end) {
            isEnd = true;
        }
    }
}
